import akshare as ak
import pandas as pd
from datetime import datetime, timedelta


def get_stock_data(code, name=None):
    """
    获取股票或ETF的数据
    :param code: 股票或ETF代码，如"601088"或"510310"
    :param name: 股票名称，仅用于显示
    :return: 数据字符串
    """
    # 判断是否为ETF (以51开头的通常是ETF)
    is_etf = code.startswith('51')
    
    # 获取实时行情
    spot_data = None
    try:
        if is_etf:
            # 尝试使用ETF专用接口
            etf_spot = ak.fund_etf_spot_em()
            spot_data = etf_spot[etf_spot["代码"] == code]
            if len(spot_data) > 0:
                spot_data = spot_data.iloc[0]
        else:
            # 普通A股
            stock_spot = ak.stock_zh_a_spot_em()
            spot_data = stock_spot[stock_spot["代码"] == code]
            if len(spot_data) > 0:
                spot_data = spot_data.iloc[0]
    except Exception as e:
        print(f"获取实时行情出错: {e}")
    
    # 如果上面的方法都失败了
    if spot_data is None or len(spot_data) == 0:
        print(f"无法获取{code}的实时行情，尝试其他API...")
        spot_data = {"最新价": "N/A", "涨跌幅": "N/A", "成交量": "N/A", "换手率": "N/A"}

    # 获取历史数据
    end_date = datetime.now().strftime("%Y%m%d")
    start_date = (datetime.now() - timedelta(days=365)).strftime("%Y%m%d")
    
    hist_data = pd.DataFrame()
    try:
        if is_etf:
            # 使用ETF历史数据接口
            hist_data = ak.fund_etf_hist_sina(symbol=f"sh{code}" if code.startswith("5") else f"sz{code}",
                                             start_date=start_date, 
                                             end_date=end_date)
            if not hist_data.empty:
                hist_data = hist_data.rename(columns={
                    "date": "date", "open": "open", "close": "close",
                    "high": "high", "low": "low", "volume": "volume",
                    "amount": "amount"
                })
        else:
            # 普通A股历史数据
            hist_data = ak.stock_zh_a_hist(
                symbol=code,
                period="daily",
                start_date=start_date,
                end_date=end_date,
                adjust="qfq"
            )
            if not hist_data.empty:
                hist_data = hist_data.rename(columns={
                    "日期": "date", "开盘": "open", "收盘": "close",
                    "最高": "high", "最低": "low", "成交量": "volume",
                    "成交额": "amount"
                })
    except Exception as e:
        print(f"获取历史数据出错: {e}")
    
    # 如果第一种方法失败，尝试其他API
    if hist_data.empty:
        try:
            print("尝试使用备用API获取数据...")
            if is_etf:
                # 尝试使用通用ETF历史数据接口
                hist_data = ak.fund_etf_hist_em(symbol=code, period="daily", 
                                              start_date=start_date, end_date=end_date)
                if not hist_data.empty:
                    hist_data = hist_data.rename(columns={
                        "日期": "date", "开盘": "open", "收盘": "close",
                        "最高": "high", "最低": "low", "成交量": "volume",
                        "成交额": "amount"
                    })
            else:
                # 尝试另一个股票API
                symbol_with_prefix = f"sh{code}" if code.startswith("6") else f"sz{code}"
                hist_data = ak.stock_zh_a_daily(symbol=symbol_with_prefix, 
                                               start_date=start_date, 
                                               end_date=end_date, 
                                               adjust="qfq")
        except Exception as e:
            print(f"备用API也获取数据失败: {e}")

    # 转换为字符串格式
    asset_name = name if name else f"{'ETF' if is_etf else '股票'}{code}"
    spot_str = f"""
[{asset_name}实时数据] 时间:{datetime.now().strftime('%Y-%m-%d %H:%M')}
最新价:{spot_data.get('最新价', 'N/A')} | 涨跌幅:{spot_data.get('涨跌幅', 'N/A')}
成交量:{spot_data.get('成交量', 'N/A')}股 | 换手率:{spot_data.get('换手率', 'N/A')}
主力净流入:{spot_data.get('主力净流入', '待更新')} 
    """.strip()

    if hist_data.empty:
        hist_str = "\n[历史数据]\n无法获取历史数据，请检查网络连接或API参数"
    else:
        hist_str = "\n[历史数据]\n" + hist_data.to_string(index=False, justify='center')

    return spot_str + hist_str


def get_shenhua_data():
    """向后兼容原有函数"""
    return get_stock_data("601088", "神华能源")


if __name__ == "__main__":
    # 测试沪深300ETF数据
    print(get_stock_data("510310", "沪深300ETF"))
    # print(get_stock_data("601088", "神华能源"))